// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

#include <nebula/compute/kernel.h>

#include <cstddef>
#include <memory>
#include <sstream>
#include <string>

#include <nebula/core/buffer.h>
#include <nebula/compute/exec.h>

#include <nebula/types/type_traits.h>
#include <nebula/bits/bit_util.h>
#include <turbo/base/checked_cast.h>
#include <nebula/util/hash_util.h>
#include <turbo/log/logging.h>
#include <turbo/base/macros.h>

namespace nebula::compute {

    using nebula::internal::hash_combine;
    static constexpr size_t kHashSeed = 0;
    // ----------------------------------------------------------------------
    // KernelContext

    turbo::Result<std::shared_ptr<ResizableBuffer>> KernelContext::allocate(int64_t nbytes) {
        return allocate_resizable_buffer(nbytes, exec_ctx_->memory_pool());
    }

    turbo::Result<std::shared_ptr<ResizableBuffer>> KernelContext::allocate_bitmap(int64_t num_bits) {
        const int64_t nbytes = bit_util::BytesForBits(num_bits);
        TURBO_MOVE_OR_RAISE(std::shared_ptr<ResizableBuffer> result,
                            allocate_resizable_buffer(nbytes, exec_ctx_->memory_pool()));
        // Since bitmaps are typically written bit by bit, we could leak uninitialized bits.
        // Make sure all memory is initialized (this also appeases Valgrind).
        std::memset(result->mutable_data(), 0, result->size());
        return result;
    }

    turbo::Status Kernel::InitAll(KernelContext *ctx, const KernelInitArgs &args,
                                  std::vector<std::unique_ptr<KernelState>> *states) {
        for (auto &state: *states) {
            TURBO_MOVE_OR_RAISE(state, args.kernel->init(ctx, args));
        }
        return turbo::OkStatus();
    }

    turbo::Result<std::unique_ptr<KernelState>> ScalarAggregateKernel::MergeAll(
            const ScalarAggregateKernel *kernel, KernelContext *ctx,
            std::vector<std::unique_ptr<KernelState>> states) {
        auto out = std::move(states.back());
        states.pop_back();
        ctx->set_state(out.get());
        for (auto &state: states) {
            TURBO_RETURN_NOT_OK(kernel->merge(ctx, std::move(*state), out.get()));
        }
        return out;
    }

    // ----------------------------------------------------------------------
    // Some basic TypeMatcher implementations

    namespace match {

        class SameTypeIdMatcher : public TypeMatcher {
        public:
            explicit SameTypeIdMatcher(Type::type accepted_id) : accepted_id_(accepted_id) {}

            bool matches(const DataType &type) const override { return type.id() == accepted_id_; }

            std::string to_string() const override {
                std::stringstream ss;
                ss << "Type::" << ::nebula::internal::to_string(accepted_id_);
                return ss.str();
            }

            bool equals(const TypeMatcher &other) const override {
                if (this == &other) {
                    return true;
                }
                auto casted = dynamic_cast<const SameTypeIdMatcher *>(&other);
                if (casted == nullptr) {
                    return false;
                }
                return this->accepted_id_ == casted->accepted_id_;
            }

        private:
            Type::type accepted_id_;
        };

        std::shared_ptr<TypeMatcher> SameTypeId(Type::type type_id) {
            return std::make_shared<SameTypeIdMatcher>(type_id);
        }

        template<typename ArrowType>
        class TimeUnitMatcher : public TypeMatcher {
            using ThisType = TimeUnitMatcher<ArrowType>;

        public:
            explicit TimeUnitMatcher(TimeUnit::type accepted_unit)
                    : accepted_unit_(accepted_unit) {}

            bool matches(const DataType &type) const override {
                if (type.id() != ArrowType::type_id) {
                    return false;
                }
                const auto &time_type = turbo::checked_cast<const ArrowType &>(type);
                return time_type.unit() == accepted_unit_;
            }

            bool equals(const TypeMatcher &other) const override {
                if (this == &other) {
                    return true;
                }
                auto casted = dynamic_cast<const ThisType *>(&other);
                if (casted == nullptr) {
                    return false;
                }
                return this->accepted_unit_ == casted->accepted_unit_;
            }

            std::string to_string() const override {
                std::stringstream ss;
                ss << ArrowType::type_name() << "(" << ::nebula::internal::to_string(accepted_unit_)
                   << ")";
                return ss.str();
            }

        private:
            TimeUnit::type accepted_unit_;
        };

        using DurationTypeUnitMatcher = TimeUnitMatcher<DurationType>;
        using Time32TypeUnitMatcher = TimeUnitMatcher<Time32Type>;
        using Time64TypeUnitMatcher = TimeUnitMatcher<Time64Type>;
        using TimestampTypeUnitMatcher = TimeUnitMatcher<TimestampType>;

        std::shared_ptr<TypeMatcher> TimestampTypeUnit(TimeUnit::type unit) {
            return std::make_shared<TimestampTypeUnitMatcher>(unit);
        }

        std::shared_ptr<TypeMatcher> Time32TypeUnit(TimeUnit::type unit) {
            return std::make_shared<Time32TypeUnitMatcher>(unit);
        }

        std::shared_ptr<TypeMatcher> Time64TypeUnit(TimeUnit::type unit) {
            return std::make_shared<Time64TypeUnitMatcher>(unit);
        }

        std::shared_ptr<TypeMatcher> DurationTypeUnit(TimeUnit::type unit) {
            return std::make_shared<DurationTypeUnitMatcher>(unit);
        }

        class IntegerMatcher : public TypeMatcher {
        public:
            IntegerMatcher() {}

            bool matches(const DataType &type) const override { return is_integer(type.id()); }

            bool equals(const TypeMatcher &other) const override {
                if (this == &other) {
                    return true;
                }
                auto casted = dynamic_cast<const IntegerMatcher *>(&other);
                return casted != nullptr;
            }

            std::string to_string() const override { return "integer"; }
        };

        std::shared_ptr<TypeMatcher> Integer() { return std::make_shared<IntegerMatcher>(); }

        class PrimitiveMatcher : public TypeMatcher {
        public:
            PrimitiveMatcher() {}

            bool matches(const DataType &type) const override { return is_primitive(type.id()); }

            bool equals(const TypeMatcher &other) const override {
                if (this == &other) {
                    return true;
                }
                auto casted = dynamic_cast<const PrimitiveMatcher *>(&other);
                return casted != nullptr;
            }

            std::string to_string() const override { return "primitive"; }
        };

        std::shared_ptr<TypeMatcher> Primitive() { return std::make_shared<PrimitiveMatcher>(); }

        class BinaryLikeMatcher : public TypeMatcher {
        public:
            BinaryLikeMatcher() {}

            bool matches(const DataType &type) const override { return is_binary_like(type.id()); }

            bool equals(const TypeMatcher &other) const override {
                if (this == &other) {
                    return true;
                }
                auto casted = dynamic_cast<const BinaryLikeMatcher *>(&other);
                return casted != nullptr;
            }

            std::string to_string() const override { return "binary-like"; }
        };

        std::shared_ptr<TypeMatcher> BinaryLike() {
            return std::make_shared<BinaryLikeMatcher>();
        }

        class LargeBinaryLikeMatcher : public TypeMatcher {
        public:
            LargeBinaryLikeMatcher() {}

            bool matches(const DataType &type) const override {
                return is_large_binary_like(type.id());
            }

            bool equals(const TypeMatcher &other) const override {
                if (this == &other) {
                    return true;
                }
                auto casted = dynamic_cast<const LargeBinaryLikeMatcher *>(&other);
                return casted != nullptr;
            }

            std::string to_string() const override { return "large-binary-like"; }
        };

        class FixedSizeBinaryLikeMatcher : public TypeMatcher {
        public:
            FixedSizeBinaryLikeMatcher() {}

            bool matches(const DataType &type) const override {
                return is_fixed_size_binary(type.id());
            }

            bool equals(const TypeMatcher &other) const override {
                if (this == &other) {
                    return true;
                }
                auto casted = dynamic_cast<const FixedSizeBinaryLikeMatcher *>(&other);
                return casted != nullptr;
            }

            std::string to_string() const override { return "fixed-size-binary-like"; }
        };

        std::shared_ptr<TypeMatcher> LargeBinaryLike() {
            return std::make_shared<LargeBinaryLikeMatcher>();
        }

        std::shared_ptr<TypeMatcher> FixedSizeBinaryLike() {
            return std::make_shared<FixedSizeBinaryLikeMatcher>();
        }

        class RunEndIntegerMatcher : public TypeMatcher {
        public:
            ~RunEndIntegerMatcher() override = default;

            bool matches(const DataType &type) const override { return is_run_end_type(type.id()); }

            bool equals(const TypeMatcher &other) const override {
                auto casted = dynamic_cast<const RunEndIntegerMatcher *>(&other);
                return casted != nullptr;
            }

            std::string to_string() const override { return "run-end-integer"; }
        };

        std::shared_ptr<TypeMatcher> RunEndInteger() {
            return std::make_shared<RunEndIntegerMatcher>();
        }

        class RunEndEncodedMatcher : public TypeMatcher {
        public:
            RunEndEncodedMatcher(std::shared_ptr<TypeMatcher> run_end_type_matcher,
                                 std::shared_ptr<TypeMatcher> value_type_matcher)
                    : run_end_type_matcher{std::move(run_end_type_matcher)},
                      value_type_matcher{std::move(value_type_matcher)} {}

            ~RunEndEncodedMatcher() override = default;

            bool matches(const DataType &type) const override {
                if (type.id() == Type::RUN_END_ENCODED) {
                    const auto &ree_type = dynamic_cast<const RunEndEncodedType &>(type);
                    // This invariant is enforced in RunEndEncodedType's constructor
                            DKCHECK(is_run_end_type(ree_type.run_end_type()->id()));
                    return run_end_type_matcher->matches(*ree_type.run_end_type()) &&
                           value_type_matcher->matches(*ree_type.get_value_type());
                }
                return false;
            }

            bool equals(const TypeMatcher &other) const override {
                if (this == &other) {
                    return true;
                }
                const auto *casted = dynamic_cast<const RunEndEncodedMatcher *>(&other);
                return casted != nullptr && value_type_matcher->equals(*casted->value_type_matcher) &&
                       run_end_type_matcher->equals(*casted->run_end_type_matcher);
            }

            std::string to_string() const override {
                return "run_end_encoded(" + run_end_type_matcher->to_string() + ", " +
                       value_type_matcher->to_string() + ")";
            };

        private:
            std::shared_ptr<TypeMatcher> run_end_type_matcher;
            std::shared_ptr<TypeMatcher> value_type_matcher;
        };

        std::shared_ptr<TypeMatcher> RunEndEncoded(
                std::shared_ptr<TypeMatcher> value_type_matcher) {
            return std::make_shared<RunEndEncodedMatcher>(RunEndInteger(),
                                                          std::move(value_type_matcher));
        }

        std::shared_ptr<TypeMatcher> RunEndEncoded(Type::type value_type_id) {
            return RunEndEncoded(SameTypeId(value_type_id));
        }

        std::shared_ptr<TypeMatcher> RunEndEncoded(
                std::shared_ptr<TypeMatcher> run_end_type_matcher,
                std::shared_ptr<TypeMatcher> value_type_matcher) {
            return std::make_shared<RunEndEncodedMatcher>(std::move(run_end_type_matcher),
                                                          std::move(value_type_matcher));
        }

    }  // namespace match

    // ----------------------------------------------------------------------
    // InputType

    size_t InputType::hash() const {
        size_t result = kHashSeed;
        hash_combine(result, static_cast<int>(kind_));
        switch (kind_) {
            case InputType::EXACT_TYPE:
                hash_combine(result, type_->hash());
                break;
            case InputType::ANY_TYPE:
            case InputType::USE_TYPE_MATCHER:
                break;
        }
        return result;
    }

    std::string InputType::to_string() const {
        std::stringstream ss;
        switch (kind_) {
            case InputType::ANY_TYPE:
                ss << "any";
                break;
            case InputType::EXACT_TYPE:
                ss << type_->to_string();
                break;
            case InputType::USE_TYPE_MATCHER: {
                ss << type_matcher_->to_string();
                break;
            }
        }
        return ss.str();
    }

    bool InputType::equals(const InputType &other) const {
        if (this == &other) {
            return true;
        }
        if (kind_ != other.kind_) {
            return false;
        }
        switch (kind_) {
            case InputType::ANY_TYPE:
                return true;
            case InputType::EXACT_TYPE:
                return type_->equals(*other.type_);
            case InputType::USE_TYPE_MATCHER:
                return type_matcher_->equals(*other.type_matcher_);
        }
        return false;
    }

    bool InputType::matches(const DataType &type) const {
        switch (kind_) {
            case InputType::EXACT_TYPE:
                return type_->equals(type);
            case InputType::USE_TYPE_MATCHER:
                return type_matcher_->matches(type);
            case InputType::ANY_TYPE:
                return true;
        }
        return false;
    }

    bool InputType::matches(const Datum &value) const {
        switch (value.kind()) {
            case Datum::NONE:
            case Datum::RECORD_BATCH:
            case Datum::TABLE:
                        DKCHECK(false) << "matches expects ARRAY, CHUNKED_ARRAY or SCALAR";
                return false;
            case Datum::ARRAY:
            case Datum::CHUNKED_ARRAY:
            case Datum::SCALAR:
                break;
        }
        return matches(*value.type());
    }

    const std::shared_ptr<DataType> &InputType::type() const {
                DKCHECK_EQ(InputType::EXACT_TYPE, kind_);
        return type_;
    }

    const TypeMatcher &InputType::type_matcher() const {
                DKCHECK_EQ(InputType::USE_TYPE_MATCHER, kind_);
        return *type_matcher_;
    }

    // ----------------------------------------------------------------------
    // OutputType

    turbo::Result<TypeHolder> OutputType::resolve(KernelContext *ctx,
                                                  const std::vector<TypeHolder> &types) const {
        switch (kind_) {
            case OutputType::FIXED:
                return type_;
            case OutputType::COMPUTED:
                break;
        }
        return resolver_(ctx, types);
    }

    const std::shared_ptr<DataType> &OutputType::type() const {
                DKCHECK_EQ(FIXED, kind_);
        return type_;
    }

    const OutputType::Resolver &OutputType::resolver() const {
                DKCHECK_EQ(COMPUTED, kind_);
        return resolver_;
    }

    std::string OutputType::to_string() const {
        switch (kind_) {
            case OutputType::FIXED:
                return type_->to_string();
            case OutputType::COMPUTED:
                break;
        }
        return "computed";
    }

    // ----------------------------------------------------------------------
    // KernelSignature

    KernelSignature::KernelSignature(std::vector<InputType> in_types, OutputType out_type,
                                     bool is_varargs)
            : in_types_(std::move(in_types)),
              out_type_(std::move(out_type)),
              is_varargs_(is_varargs),
              hash_code_(0) {
                DKCHECK(!is_varargs || (is_varargs && (in_types_.size() >= 1)));
    }

    std::shared_ptr<KernelSignature> KernelSignature::create(std::vector<InputType> in_types,
                                                           OutputType out_type,
                                                           bool is_varargs) {
        return std::make_shared<KernelSignature>(std::move(in_types), std::move(out_type),
                                                 is_varargs);
    }

    bool KernelSignature::equals(const KernelSignature &other) const {
        if (is_varargs_ != other.is_varargs_) {
            return false;
        }
        if (in_types_.size() != other.in_types_.size()) {
            return false;
        }
        for (size_t i = 0; i < in_types_.size(); ++i) {
            if (!in_types_[i].equals(other.in_types_[i])) {
                return false;
            }
        }
        return true;
    }

    bool KernelSignature::matches_inputs(const std::vector<TypeHolder> &types) const {
        if (is_varargs_) {
            for (size_t i = 0; i < types.size(); ++i) {
                if (!in_types_[std::min(i, in_types_.size() - 1)].matches(*types[i])) {
                    return false;
                }
            }
        } else {
            if (types.size() != in_types_.size()) {
                return false;
            }
            for (size_t i = 0; i < in_types_.size(); ++i) {
                if (!in_types_[i].matches(*types[i])) {
                    return false;
                }
            }
        }
        return true;
    }

    size_t KernelSignature::hash() const {
        if (hash_code_ != 0) {
            return hash_code_;
        }
        size_t result = kHashSeed;
        for (const auto &in_type: in_types_) {
            hash_combine(result, in_type.hash());
        }
        hash_code_ = result;
        return result;
    }

    std::string KernelSignature::to_string() const {
        std::stringstream ss;

        if (is_varargs_) {
            ss << "varargs[";
        } else {
            ss << "(";
        }
        for (size_t i = 0; i < in_types_.size(); ++i) {
            if (i > 0) {
                ss << ", ";
            }
            ss << in_types_[i].to_string();
        }
        if (is_varargs_) {
            ss << "*]";
        } else {
            ss << ")";
        }
        ss << " -> " << out_type_.to_string();
        return ss.str();
    }

}  // namespace nebula::compute
